/*=========================================================================

Program:   Insight Segmentation & Registration Toolkit
Module:    $RCSfile: TotalVariationDenoisingImageFilter.cxx,v $
Language:  C++
Date:      $Date: 2007-12-08 18:26:26 $
Version:   $Revision: x $

Copyright (c) Insight Software Consortium. All rights reserved.
See ITKCopyright.txt or http://www.itk.org/HTML/Copyright.htm for details.

This software is distributed WITHOUT ANY WARRANTY; without even 
the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR 
PURPOSE.  See the above copyright notices for more information.

=========================================================================*/


// 
// This example shows how to use the TotalVariationDenoisingImageFilter
// to denoise images
//
//     Run the example with the following args
//       test.nrrd 20 0.1 2 out.nrrd

// itk includes
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkTimeProbesCollectorBase.h"

#include "itkTotalVariationDenoisingImageFilter.h"

// vector image typedefs
#define VECTOR_LENGTH 65

typedef itk::Vector<float,VECTOR_LENGTH> 
        VectorPixelType;
typedef itk::Image<VectorPixelType, 3>
        ImageType;
typedef itk::ImageRegionIterator<ImageType> 
        IteratorType;

int main(int argc, char **argv)
{

  if(argc < 3)
  {
    std::cerr << "Digital TV Image Denoising"
      << std::endl << std::endl;

    std::cerr << "Usage: " << argv[0] << " FileName(.nhdr/.nrrd)"
      << " #Iterations Lambda [#Threads] [OutFileName]" 
      << std::endl << std::endl;

    std::cerr << "\tExample args: test.nrrd 20 10 4 out.nrrd" << std::endl;
    std::cerr << "\tExample args: test 20 10.nrrd" << std::endl;
    return EXIT_FAILURE;
  }

  ////////////////////////
  // typedefs
  typedef itk::VectorImage<float, 3>     
    VarVecImgType;
  typedef itk::TotalVariationDenoisingImageFilter<ImageType,ImageType>
    FilterType;
  typedef itk::ImageFileReader< VarVecImgType > 
    ReaderType;
  typedef itk::ImageFileWriter< VarVecImgType >     
    WriterType;

  ////////////////////////
  // Read the input file
  ReaderType::Pointer reader = ReaderType::New();
  reader->SetFileName(argv[1]);

  try
  {
    reader->UpdateLargestPossibleRegion();
  }
  catch (itk::ExceptionObject& e)
  {
    std::cerr << "Exception detected while reading " << argv[1];
    std::cerr << " : "  << e.GetDescription();
    return EXIT_FAILURE;
  }

  VarVecImgType::Pointer image = reader->GetOutput();

  ////////////////////////
  // convert image to fixed length vector image
  ImageType::Pointer img = ImageType::New();
  img->SetSpacing( image->GetSpacing() );
  img->SetOrigin( image->GetOrigin() );
  img->SetDirection( image->GetDirection());
  img->SetLargestPossibleRegion( 
    image->GetLargestPossibleRegion());
  img->SetBufferedRegion( 
    image->GetLargestPossibleRegion() );
  img->Allocate();

  itk::ImageRegionIterator<ImageType> ot (
    img, img->GetLargestPossibleRegion() );
  ot = ot.Begin();

  itk::ImageRegionIterator<VarVecImgType> it (
    image, 
    image->GetLargestPossibleRegion() );
  it = it.Begin();

  for (it = it.Begin(); !it.IsAtEnd(); ++it)
  {
    VarVecImgType::PixelType varvec = it.Get();
    VectorPixelType vec(varvec.GetDataPointer());
    ot.Set(vec);
    ++ot;
  }

  ////////////////////////
  // init and run denoising
  FilterType::Pointer filter = 
    FilterType::New();

  filter->SetInput(img);
  filter->SetNumberIterations(atof(argv[2]));
  filter->SetLambda(atof(argv[3]));

  if(argc > 4)
    filter->SetNumberOfThreads(atoi(argv[4]));

  std::cout << std::endl << "This filter is using " << 
    filter->GetNumberOfThreads() << " threads " << std::endl;

  // start timing
  itk::TimeProbesCollectorBase collector;
  collector.Start( "TotalTime" );

  // run
  filter->Update();

  // stop and report timing
  collector.Stop( "TotalTime" );
  collector.Report();

  ImageType::Pointer resultImage = filter->GetOutput();

  ////////////////////////
  // convert result to variable length vector image that can be written to file
  // (simple image with pixeltype "fixed vector" not supported by writer)
  VarVecImgType::Pointer vecImg = VarVecImgType::New();
  vecImg->SetSpacing( resultImage->GetSpacing() );
  vecImg->SetOrigin( resultImage->GetOrigin() );
  vecImg->SetDirection( resultImage->GetDirection());
  vecImg->SetLargestPossibleRegion( 
    resultImage->GetLargestPossibleRegion());
  vecImg->SetBufferedRegion( 
    resultImage->GetLargestPossibleRegion() );
  vecImg->SetVectorLength(VECTOR_LENGTH);
  vecImg->Allocate();

  itk::ImageRegionIterator<VarVecImgType> ot2 (
    vecImg, vecImg->GetLargestPossibleRegion() );
  ot2 = ot2.Begin();

  itk::ImageRegionIterator<ImageType> it2 (
    resultImage, 
    resultImage->GetLargestPossibleRegion() );
  it2 = it2.Begin();

  for (it2 = it2.Begin(); !it2.IsAtEnd(); ++it2)
  {
    VectorPixelType vec = it2.Get();
    VarVecImgType::PixelType varvec(vec.GetDataPointer(), VECTOR_LENGTH);
    ot2.Set(varvec);
    ++ot2;
  }

  vecImg->SetMetaDataDictionary(reader->GetOutput()->GetMetaDataDictionary());

  ////////////////////////
  // Write out the image
  WriterType::Pointer writer = WriterType::New();
  writer->SetInput(vecImg);
  writer->UseInputMetaDataDictionaryOn();
  ::itk::OStringStream baseName;
  if(argc > 5)
  {
    baseName << argv[5];
  }
  else
  {
    baseName << argv[1] << ".result.nrrd";
  }

  try
  {
    writer->SetFileName( baseName.str().c_str() );
    writer->Update();
  }
  catch (itk::ExceptionObject& e)
  {
    std::cerr << "Error during write of " << baseName.str();
    std::cerr << " : "  << e.GetDescription();
    return EXIT_FAILURE;
  }

  return EXIT_SUCCESS;
}


